import parser
import os
from .base import Base

class PaddleDynamic2ONNX(Base):
    def __init__(self, success_list, failed_list):
        print('PaddleDynamic2ONNX Init--------------------------')
        self.tool_path = os.path.join(parser.get_value('maca_converter_root'), parser.get_value('maca_converter_path'))
        self.model_path = os.path.join(parser.get_value('model_root'), parser.get_value('paddle2onnx_dynamic_model_path'))
        self.model_def_file = os.path.join(parser.get_value('model_root'), parser.get_value('paddle2onnx_dynamic_model_def_file'))
        self.model_class_name = parser.get_value('paddle2onnx_dynamic_model_class_name')
        self.input_shape = parser.get_value('paddle2onnx_dynamic_model_input_shape')
        self.model_weights_file = os.path.join(parser.get_value('model_root'), parser.get_value('paddle2onnx_dynamic_model_weights_file'))
        self.success_list = success_list
        self.failed_list = failed_list
        self.input_npy = os.path.join(parser.get_value('model_root'), parser.get_value('paddle2onnx_dynamic_model_input_npy'))
        self.output_npy = os.path.join(parser.get_value('model_root'), parser.get_value('paddle2onnx_dynamic_model_output_npy'))


    def test(self):
        cwd = os.getcwd()
        output = os.path.join(cwd, 'tmp.onnx')

        cmd = 'python ' + 'model_convert.py' + ' --model_path ' + self.model_path + ' --model_type paddle ' + '--output ' + output + \
                ' --model_def_file ' + self.model_def_file + ' --model_class_name ' + self.model_class_name + \
                ' --input_shape ' + self.input_shape + ' --model_weights_file ' + self.model_weights_file
    
        print('test_paddle_dynamic_2_onnx: ', cmd)

        f = self.model_def_file.split('/')[-1]
        f = os.path.join(self.tool_path, f)
        self.mx_copy(self.model_def_file, f)

        self.test_model2onnx(self.tool_path, 'paddle_dynamic2onnx', cmd, output, self.success_list, self.failed_list, self.input_npy, self.output_npy)